import gradio as gr
#from document2llm.md.marker.marker_pdf2markdown import PDFConverter
from  .marker_pdf2markdown import PDFConverter
import os

# 从环境变量获取端口，如果没有设置，则使用默认值7860
DEFAULT_PORT = int(os.environ.get('GRADIO_SERVER_PORT', 34922 ))

examples = [
    ["/Users/Shared/work_cache/input_cache/", "/Users/Shared/work_cache/output_cache", "目录", 1, 1000],
    ["~/Downloads/1-dev/input/", "", "目录", 1, 1000],
    ["~/Downloads/1-dev/input/test.pdf", "", "文件", 1, 1000],
    ["path/to/pdf_folder", "", "目录", 2, 500],
    ["~/Documents/important.pdf", "/Users/username/Desktop/output", "文件", 1, 100],
]

home = os.path.expanduser("~")

def get_default_output_folder():
    return os.path.join(home, "Downloads")

def process_pdf(input_file, input_path, output_folder, is_directory_input, batch_multiplier, max_pages):
    if input_file:
        input_path = input_file.name
        is_directory_input = 0  # 设置为文件类型
    elif input_path:
        input_path = input_path.replace('~', home)
    else:
        return "请选择输入文件或输入路径"

    #
    input_path = input_path.strip()

    if not output_folder.strip():
        output_folder = get_default_output_folder()

    #
    output_folder = output_folder.strip()

    converter = PDFConverter(
        max_pages=max_pages,
        batch_multiplier=batch_multiplier
    )

    try:
        if is_directory_input == 1:  # 目录
            converter.process_directory(input_path, output_folder)
            return f"处理完成。输出目录: {output_folder}"
        else:  # 文件
            subfolder_path = converter.process_file(input_path, output_folder)
            if subfolder_path:
                return f"处理完成。输出子文件夹: {subfolder_path}"
            else:
                return "文件处理失败"
    except Exception as e:
        return f"发生错误: {str(e)}"

def update_input_type(file):
    if file is not None:
        return 0  # 0 代表 "文件"
    return gr.update()  # 不更改当前值

# 创建Gradio界面
with gr.Blocks() as iface:
    gr.Markdown("# PDF转Markdown工具")
    gr.Markdown("将PDF文件或目录中的PDF文件转换为Markdown格式。")
    
    with gr.Row():
        input_file = gr.File(label="选择PDF文件", file_types=[".pdf"])
        input_path = gr.Textbox(label="或输入PDF文件或目录路径")
    
    with gr.Row():
        output_folder = gr.Textbox(label="输出文件夹路径", placeholder=f"默认: {get_default_output_folder()}")

    with gr.Row():
        is_directory_input = gr.Radio(
            ["文件", "目录"], 
            label="输入类型", 
            value="文件",
            type="index"
        )
        batch_multiplier = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="批处理乘数")
        max_pages = gr.Number(value=1000, label="最大页数")
    
    process_button = gr.Button("开始处理")
    result = gr.Textbox(label="处理结果")
    
    input_file.change(
        fn=update_input_type,
        inputs=[input_file],
        outputs=[is_directory_input]
    )
    
    process_button.click(
        fn=process_pdf,
        inputs=[input_file, input_path, output_folder, is_directory_input, batch_multiplier, max_pages],
        outputs=result
    )

    gr.Markdown("## 使用示例")
    gr.Examples(
        examples=examples,
        inputs=[input_path, output_folder, is_directory_input, batch_multiplier, max_pages],
    )

# 运行Gradio应用
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="PDF to Markdown Converter")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Port to run the Gradio app on")
    args = parser.parse_args()

    print(f"Starting Gradio app on port {args.port}")
    iface.launch(server_name="0.0.0.0", server_port=args.port)